import torch
from Batch import nopeak_mask
import torch.nn.functional as F
import math


def init_vars(src, model, SRC, TRG, opt):
    init_tok = TRG.vocab.stoi['<sos>']
    src_mask = (src != SRC.vocab.stoi['<pad>']).unsqueeze(-2)
    e_output = model.encoder(src, src_mask)

    outputs = torch.LongTensor([[init_tok]])
    if opt.device == 0:
        outputs = outputs.cuda()

    trg_mask = nopeak_mask(1, opt)  # 做掩码？

    out = model.out(model.decoder(outputs,e_output, src_mask, trg_mask))
    out = F.softmax(out, dim=-1)
    probs, ix = out[:, -1].data.topk(opt.k)
    log_scores = torch.Tensor([math.log(prob) for prob in probs.data[0]]).unsqueeze(0)

    outputs = torch.zeros(opt.k, opt.max_len).long()
    if opt.device == 0:
        outputs = outputs.cuda()
    outputs[:, 0] = init_tok
    outputs[:, 1] = ix[0]

    e_outputs = torch.zeros(opt.k, e_output.size(-2), e_output.size(-1))
    if opt.device == 0:
        e_outputs = e_outputs.cuda()
    e_outputs[:, :] = e_output[0]
    # print("outputs, e_outputs, log_scores",outputs, e_outputs, log_scores)

    return outputs, e_outputs, log_scores


def k_best_outputs(outputs, out, log_scores, i, k):
    probs, ix = out[:, -1].data.topk(k)
    log_probs = torch.Tensor([math.log(p) for p in probs.data.view(-1)]).view(k, -1) + log_scores.transpose(0, 1)
    k_probs, k_ix = log_probs.view(-1).topk(k)

    row = k_ix // k
    col = k_ix % k

    outputs[:, :i] = outputs[row, :i]
    outputs[:, i] = ix[row, col]

    log_scores = k_probs.unsqueeze(0)

    return outputs, log_scores


def beam_search(src, model, SRC, TRG, opt):  # 真正的预测代码
    outputs, e_outputs, log_scores = init_vars(src, model, SRC, TRG, opt)

    eos_tok = TRG.vocab.stoi['<eos>']
    src_mask = (src != SRC.vocab.stoi['<pad>']).unsqueeze(-2)

    ind = None

    trg_mask = nopeak_mask(6, opt)  # 对角下面为True 上面为False
    out = model.out(model.decoder(outputs[:, :6], e_outputs, src_mask, trg_mask))
    out = F.softmax(out, dim=-1)
    outputs, log_scores = k_best_outputs(outputs, out, log_scores, 6, opt.k)


    for i in range(2, opt.max_len):
        trg_mask = nopeak_mask(i, opt)  # 对角下面为True 上面为False
        out = model.out(model.decoder(outputs[:, :i], e_outputs, src_mask, trg_mask))
        out = F.softmax(out, dim=-1)
        outputs, log_scores = k_best_outputs(outputs, out, log_scores, i, opt.k)
        # print(' '.join([TRG.vocab.itos[tok] for tok in outputs[0][1:6]]))
    result=[]
    for i in range(opt.k):
        length = (outputs[i] == eos_tok).nonzero()[0]
        result.append(' '.join([TRG.vocab.itos[tok] for tok in outputs[i][1:length]]))
    # else:
    #     length = (outputs[ind] == eos_tok).nonzero()[0]
    #     return ' '.join([TRG.vocab.itos[tok] for tok in outputs[ind][1:length]])
    return result